#!/usr/bin/env python3
"""
检查训练好的模型权重结构
"""

import torch
from pathlib import Path

model_path = Path(r"D:\VSCodeProjects\UC_System\uc_model\outputs\standard_training\best_model.pth")

print(f"检查模型文件: {model_path}")
print(f"文件大小: {model_path.stat().st_size / 1024 / 1024:.1f} MB")

# 加载权重
checkpoint = torch.load(model_path, map_location='cpu')

print("\n权重文件结构:")
if isinstance(checkpoint, dict):
    print(f"- 键的数量: {len(checkpoint)}")
    print("- 主要键:", list(checkpoint.keys())[:5])

    if 'model_state_dict' in checkpoint:
        state_dict = checkpoint['model_state_dict']
        print("\n模型权重包含的主要组件:")
        keys = list(state_dict.keys())

        # 统计不同前缀的键
        prefixes = {}
        for key in keys:
            prefix = key.split('.')[0]
            if prefix not in prefixes:
                prefixes[prefix] = 0
            prefixes[prefix] += 1

        for prefix, count in sorted(prefixes.items()):
            print(f"- {prefix}: {count} 个参数")

        # 检查是否有多模态相关的键
        has_image_branch = any('image_branch' in k for k in keys)
        has_text_branch = any('text_branch' in k for k in keys)
        has_fusion = any('fusion' in k for k in keys)

        print(f"\n模型特征:")
        print(f"- 有图像分支: {has_image_branch}")
        print(f"- 有文本分支: {has_text_branch}")
        print(f"- 有融合层: {has_fusion}")

    elif 'state_dict' in checkpoint:
        print("\n包含 'state_dict' 键")
    else:
        # 直接是state_dict
        print("\n直接包含模型权重（state_dict）")
        keys = list(checkpoint.keys())
        print(f"权重数量: {len(keys)}")
        print("前10个键:", keys[:10])

else:
    print("权重文件不是字典格式")